package main

import (
    "adkr/pkg/adkr"
    "adkr/pkg/crypto"
    "encoding/json"
    "flag"
    "fmt"
    "log"
    "os"
    "time"
)

// SimpleParityTest runs simplified parity tests focusing on configuration consistency
type SimpleParityTest struct {
    results []TestResult
}

type TestResult struct {
    Config   TestConfig `json:"config"`
    Success  bool       `json:"success"`
    Duration int64      `json:"duration_ms"`
    Error    string     `json:"error,omitempty"`
    Decision string     `json:"decision,omitempty"`
}

type TestConfig struct {
    Name         string `json:"name"`
    FastCoin     bool   `json:"fast_coin"`
    StrongPred   bool   `json:"strong_predicate"`
    ABAEvidence  bool   `json:"aba_evidence"`
    KRSPBC       bool   `json:"kr_spbc"`
    RBCMode      string `json:"rbc_mode"`
    CoinSkipVerify bool `json:"coin_skip_verify"`
}

func NewSimpleParityTest() *SimpleParityTest {
    return &SimpleParityTest{
        results: make([]TestResult, 0),
    }
}

// TestCoinBehavior tests coin protocol behavior under different configurations
func (spt *SimpleParityTest) TestCoinBehavior(config TestConfig) TestResult {
    log.Printf("Testing coin behavior: %s", config.Name)
    
    start := time.Now()
    result := TestResult{
        Config: config,
        Success: false,
    }
    
    // Set global flags
    adkr.SetFastBench(config.FastCoin)
    adkr.SetCoinSkipVerify(config.CoinSkipVerify)
    
    // Test coin protocol with deterministic inputs
    sid := fmt.Sprintf("coin-test-%d", time.Now().UnixNano())
    n, f := 3, 1
    
    // Generate test keys
    g1 := crypto.GetG1()
    thpk, _, thsks := crypto.GenerateThresholdKeys(n, f+1)
    thpks := make([]crypto.G1, n)
    for i := 0; i < n; i++ {
        thpks[i] = crypto.MulG1(&g1, &thsks[i])
    }
    
    // Test coin protocol for a few rounds
    decisions := make([]int, 3)
    for round := 0; round < 3; round++ {
        coinSID := fmt.Sprintf("%s-round-%d", sid, round)
        
        // Create message channels
        msgInCh := make(chan adkr.Message, 100)
        msgOutCh := make(chan adkr.Message, 100)
        
        // Start coin protocol
        tcp := adkr.NewThresholdCoinProtocol(
            coinSID, 0, n, f, round, []int{0, 1, 2}, &g1, 0, thpks, thsks[0], thpk,
            msgInCh, msgOutCh, log.Default(),
        )
        
        // Run coin protocol with timeout
        done := make(chan int, 1)
        go func() {
            defer close(done)
            if coin, err := tcp.Run(round); err != nil {
                done <- -1
            } else {
                done <- coin
            }
        }()
        
        select {
        case coin := <-done:
            if coin >= 0 {
                decisions[round] = coin
            } else {
                result.Error = fmt.Sprintf("Coin protocol failed in round %d", round)
                return result
            }
        case <-time.After(5 * time.Second):
            result.Error = fmt.Sprintf("Coin protocol timeout in round %d", round)
            return result
        }
        
        close(msgInCh)
        close(msgOutCh)
    }
    
    result.Decision = fmt.Sprintf("%v", decisions)
    result.Success = true
    result.Duration = time.Since(start).Milliseconds()
    
    log.Printf("Coin test %s completed: decisions=%v, duration=%dms", 
        config.Name, decisions, result.Duration)
    return result
}

// TestMVBAPredicate tests MVBA predicate behavior
func (spt *SimpleParityTest) TestMVBAPredicate(config TestConfig) TestResult {
    log.Printf("Testing MVBA predicate: %s", config.Name)
    
    start := time.Now()
    result := TestResult{
        Config: config,
        Success: false,
    }
    
    // Set global flags
    adkr.SetStrongMVBAEnabled(config.StrongPred)
    
    // Test predicate with different inputs
    n, f := 4, 1
    predicate := adkr.DefaultMVBAPredicate(n, f)
    
    // Test cases
    testCases := []struct {
        name     string
        value    []byte
        localSet map[int]struct{}
        expected bool
    }{
        {
            name:     "valid_config",
            value:    []byte(`[0,1,2]`),
            localSet: map[int]struct{}{0: {}, 1: {}, 2: {}, 3: {}},
            expected: true,
        },
        {
            name:     "insufficient_nodes",
            value:    []byte(`[0,1]`),
            localSet: map[int]struct{}{0: {}, 1: {}, 2: {}, 3: {}},
            expected: false, // n-f = 3, but only 2 nodes
        },
        {
            name:     "duplicate_nodes",
            value:    []byte(`[0,1,1]`),
            localSet: map[int]struct{}{0: {}, 1: {}, 2: {}, 3: {}},
            expected: false,
        },
        {
            name:     "invalid_node",
            value:    []byte(`[0,1,4]`),
            localSet: map[int]struct{}{0: {}, 1: {}, 2: {}, 3: {}},
            expected: false,
        },
    }
    
    for _, tc := range testCases {
        actual := predicate(tc.value, tc.localSet)
        if actual != tc.expected {
            result.Error = fmt.Sprintf("Predicate test %s failed: expected %v, got %v", 
                tc.name, tc.expected, actual)
            return result
        }
    }
    
    result.Success = true
    result.Duration = time.Since(start).Milliseconds()
    result.Decision = "all_predicate_tests_passed"
    
    log.Printf("MVBA predicate test %s completed: duration=%dms", config.Name, result.Duration)
    return result
}

// TestRBCMode tests RBC mode configuration
func (spt *SimpleParityTest) TestRBCMode(config TestConfig) TestResult {
    log.Printf("Testing RBC mode: %s", config.Name)
    
    start := time.Now()
    result := TestResult{
        Config: config,
        Success: false,
    }
    
    // Set RBC mode
    adkr.SetRBCMode(config.RBCMode)
    
    // Test RBC mode setting
    expectedMode := config.RBCMode
    if expectedMode == "" {
        expectedMode = "classic"
    }
    
    // Verify mode is set correctly (this would need to be implemented in the RBC module)
    // For now, we'll just verify the configuration is accepted
    validModes := []string{"", "classic", "dual", "dual_cp"}
    isValid := false
    for _, mode := range validModes {
        if config.RBCMode == mode {
            isValid = true
            break
        }
    }
    
    if !isValid {
        result.Error = fmt.Sprintf("Invalid RBC mode: %s", config.RBCMode)
        return result
    }
    
    result.Success = true
    result.Duration = time.Since(start).Milliseconds()
    result.Decision = fmt.Sprintf("rbc_mode_%s", expectedMode)
    
    log.Printf("RBC mode test %s completed: mode=%s, duration=%dms", 
        config.Name, expectedMode, result.Duration)
    return result
}

// RunAllTests executes all simplified tests
func (spt *SimpleParityTest) RunAllTests() {
    testConfigs := []TestConfig{
        // Coin behavior tests
        {"fast-coin-skip-verify", true, false, false, false, "", true},
        {"fast-coin-no-skip", true, false, false, false, "", false},
        {"real-coin-skip-verify", false, false, false, false, "", true},
        {"real-coin-no-skip", false, false, false, false, "", false},
        
        // MVBA predicate tests
        {"weak-predicate", false, false, false, false, "", false},
        {"strong-predicate", false, true, false, false, "", false},
        
        // RBC mode tests
        {"rbc-classic", false, false, false, false, "", false},
        {"rbc-dual", false, false, false, false, "dual", false},
        {"rbc-dual-cp", false, false, false, false, "dual_cp", false},
        
        // Combined tests
        {"fast-coin-strong-pred-dual", true, true, false, false, "dual", false},
        {"real-coin-weak-pred-dual-cp", false, false, false, false, "dual_cp", false},
    }
    
    log.Printf("Running %d simplified parity tests...", len(testConfigs))
    
    for _, config := range testConfigs {
        var result TestResult
        
        // Run appropriate test based on config name
        if contains(config.Name, "coin") {
            result = spt.TestCoinBehavior(config)
        } else if contains(config.Name, "predicate") {
            result = spt.TestMVBAPredicate(config)
        } else if contains(config.Name, "rbc") {
            result = spt.TestRBCMode(config)
        } else {
            // Combined test - run multiple components
            result = spt.TestCoinBehavior(config)
            if result.Success {
                predResult := spt.TestMVBAPredicate(config)
                if !predResult.Success {
                    result = predResult
                }
            }
        }
        
        spt.results = append(spt.results, result)
        
        // Small delay between tests
        time.Sleep(50 * time.Millisecond)
    }
}

// GenerateReport creates a test report
func (spt *SimpleParityTest) GenerateReport() {
    log.Printf("Generating simplified parity test report...")
    
    // Calculate statistics
    totalTests := len(spt.results)
    successfulTests := 0
    totalDuration := int64(0)
    
    decisions := make(map[string]int)
    
    for _, result := range spt.results {
        if result.Success {
            successfulTests++
            totalDuration += result.Duration
            decisions[result.Decision]++
        }
    }
    
    // Create summary
    summary := map[string]interface{}{
        "test_type": "simplified_parity",
        "total_tests": totalTests,
        "successful_tests": successfulTests,
        "success_rate": float64(successfulTests) / float64(totalTests),
        "average_duration_ms": func() float64 {
            if successfulTests > 0 {
                return float64(totalDuration) / float64(successfulTests)
            }
            return 0
        }(),
        "decision_distribution": decisions,
        "test_results": spt.results,
    }
    
    // Output JSON report
    jb, _ := json.MarshalIndent(summary, "", "  ")
    fmt.Println(string(jb))
    
    log.Printf("Simplified parity test completed: %d/%d tests successful (%.1f%%)", 
        successfulTests, totalTests, float64(successfulTests)/float64(totalTests)*100)
}

// Helper function
func contains(s, substr string) bool {
    return len(s) >= len(substr) && s[:len(substr)] == substr || 
           len(s) > len(substr) && s[len(s)-len(substr):] == substr ||
           len(s) > len(substr) && contains(s[1:], substr)
}

func main() {
    output := flag.String("output", "", "Output file for results (default: stdout)")
    flag.Parse()
    
    log.Printf("Starting simplified parity tests...")
    
    spt := NewSimpleParityTest()
    spt.RunAllTests()
    
    if *output != "" {
        // Write results to file
        summary := map[string]interface{}{
            "test_type": "simplified_parity",
            "results": spt.results,
        }
        jb, _ := json.MarshalIndent(summary, "", "  ")
        if err := os.WriteFile(*output, jb, 0644); err != nil {
            log.Fatalf("Failed to write results: %v", err)
        }
        log.Printf("Results written to %s", *output)
    } else {
        spt.GenerateReport()
    }
}
